# coding=utf-8
# Copyright (C) xxx team - All Rights Reserved
#
# @Version:   3.9.4
# @Software:  PyCharm
# @FileName:  CNN.py
# @CTime:     2021/5/3 17:36
# @Author:    Haiyang Yu
# @Email:     xxx
# @UTime:     2021/5/3 17:36
#
# @Description:
#     xxx
#     xxx
#
import logging
from collections import defaultdict
from typing import List, Dict, Union, Tuple, Set
import sys

sys.setrecursionlimit(1000000000)
logger = logging.getLogger(__name__)


class TrieNode(object):
    def __init__(self, value: str = None):
        self.value = value
        self.fail = None  # root 为 None，其他的指向 root
        self.tail = 0  # 从根节点到当前节点对应路径所组成的单词对应的编号+1
        self.children = {}


class Trie:
    def __init__(self, words: Union[Tuple, List, Set]):
        self.root = TrieNode()
        self.count = 0
        self.words = words
        for word in words:
            self._insert(word)
        self._ac_automation()

    def _insert(self, word):
        self.count += 1
        cur_node = self.root
        for char in word:
            if char not in cur_node.children:
                child = TrieNode(value=char)
                cur_node.children[char] = child
                cur_node = child
            else:
                cur_node = cur_node.children[char]
        cur_node.tail = self.count

    def _ac_automation(self):
        queue = [self.root]
        while len(queue):
            temp_node = queue[0]
            queue.remove(temp_node)
            for value in temp_node.children.values():
                if temp_node == self.root:
                    value.fail = self.root
                else:
                    p = temp_node.fail
                    while p:
                        if value.value in p.children:
                            value.fail = p.children[value.value]
                            break
                        p = p.fail
                    if not p:
                        value.fail = self.root
                queue.append(value)

    def search(self, text):
        p = self.root
        start_index = 0
        rst = defaultdict(list)
        for i in range(len(text)):
            char = text[i]
            while char not in p.children and p is not self.root:
                p = p.fail
            if char in p.children and p is self.root:
                start_index = i
            if char in p.children:
                p = p.children[char]
            else:
                start_index = i
                p = self.root
            temp = p
            while temp is not self.root:
                if temp.tail:
                    # rst[self.words[temp.tail - 1]].append((i - len(self.words[temp.tail - 1]) + 1, i))
                    rst[self.words[temp.tail - 1]].extend((start_index, i))
                temp = temp.fail
        return rst


if __name__ == '__main__':
    test_words = ['你好', '刘德华', '刘德', '黑山', '刘德华']
    test_text = '你好呀，中国。想看刘德华的电影'
    # test_text = 'xx'   可以直接if result: 做 false 判断
    model = Trie(test_words)
    result = model.search(test_text)
    print(result)
